package edu.pitt.cs.zhangfan.sa.core;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.Iterator;

import org.apache.poi.xssf.usermodel.XSSFCell;
import org.apache.poi.xssf.usermodel.XSSFRow;
import org.apache.poi.xssf.usermodel.XSSFSheet;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;

import edu.pitt.cs.zhangfan.sa.io.ExcelReader;
import edu.pitt.cs.zhangfan.sa.model.DocumentPair;
import edu.pitt.cs.zhangfan.sa.model.Sentence;

/**
 * Automatically aligns sentence index
 * 
 * @author zhangfan
 * 
 */
public class AutomaticAnnotator {
	// Set up the default value for later read, seems to be useless
	public void writeDefault(String file) throws IOException {
		FileInputStream fis = new FileInputStream(new File(file));
		XSSFWorkbook xwb = new XSSFWorkbook(fis);
		XSSFSheet sheet = xwb.getSheetAt(1);
		for (int i = sheet.getFirstRowNum(); i < sheet
				.getPhysicalNumberOfRows(); i++) {
			XSSFRow row = sheet.getRow(i);
			if (row == null)
				break;

			XSSFCell cell = row.getCell(row.getFirstCellNum() + 1);
			if (cell == null)
				cell = row.createCell(row.getFirstCellNum() + 1);
			cell.setCellValue(1);
		}
		fis.close();
		FileOutputStream fileOut = new FileOutputStream(file);
		xwb.write(fileOut);
		fileOut.flush();
		fileOut.close();

	}
	
	public void writeXSL(XSSFWorkbook xwb, int lineNo, String val) {
		XSSFSheet sheet = xwb.getSheetAt(1);
		XSSFRow row = sheet.getRow(lineNo);
		XSSFCell cell = row.getCell(row.getFirstCellNum() + 1);
		if (cell == null)
			cell = row.createCell(row.getFirstCellNum() + 1);
		cell.setCellValue(val);
	}

	public Sentence createMergedSentence(Sentence a, Sentence b) {
		Sentence sen = new Sentence();
		sen.setContent(a.getContent() + " " + b.getContent());
		return sen;
	}

	// get the segments
	public ArrayList<ArrayList<Integer>> getConsistentArrays(
			ArrayList<Integer> arrs) {
		ArrayList<ArrayList<Integer>> segments = new ArrayList<ArrayList<Integer>>();
		ArrayList<Integer> tmp = new ArrayList<Integer>();
		for (int i = 0; i < arrs.size(); i++) {
			int currIndex = arrs.get(i);
			// consistent
			if (tmp.size() == 0 || tmp.get(tmp.size() - 1) + 1 == currIndex
					|| i == 0) {
				tmp.add(currIndex);
			} else {
				ArrayList<Integer> addList = new ArrayList<Integer>();
				for (int j = 0; j < tmp.size(); j++) {
					addList.add(tmp.get(j));
				}
				segments.add(addList);
				tmp.clear();
				tmp.add(currIndex);
			}
		}

		if (tmp.size() > 0) {
			ArrayList<Integer> addList = new ArrayList<Integer>();
			for (int j = 0; j < tmp.size(); j++) {
				addList.add(tmp.get(j));
			}
			segments.add(addList);
		}
		return segments;
	}

	public double getProb(double[][] prob, ArrayList<Integer> sents,
			DocumentPair pair, int lineNo, LRAligner aligner, int option)
			throws Exception {
		if (sents.size() == 1) {
			return prob[sents.get(0)][lineNo];
		} else {
			Sentence merged = new Sentence();
			for (int i = 0; i < sents.size(); i++) {
				merged = createMergedSentence(merged, pair.getSrc()
						.getSentences().get(sents.get(i)));
			}
			return aligner.classifySingle(merged, pair.getModified()
					.getSentences().get(lineNo), option);
		}
	}

	//Spliting an array into several segments
	public ArrayList<ArrayList<Integer>> splitsArray(ArrayList<Integer> arr) {
		ArrayList<ArrayList<Integer>> segments = new ArrayList<ArrayList<Integer>>();
		for(int i = 0;i<arr.size();i++) {
			for(int j = 1;j<=arr.size()-i;j++) {
				ArrayList<Integer> segment = new ArrayList<Integer>();
				for(int k = 0;k<j;k++) {
					segment.add(arr.get(i+k));
				}
				segments.add(segment);
			}
		}
		return segments;
	}
	
	/*
	 * Code still needs to be modified
	 */
	public ArrayList<Integer> getAlignedSentences(double[][] prob, int lineNo,
			DocumentPair pair, LRAligner aligner, int option) throws Exception {
		ArrayList<Integer> candidateMatches = new ArrayList<Integer>();
		for (int i = 0; i < prob.length; i++) {
			if (prob[i][lineNo] > 0.5) {
				candidateMatches.add(i);
			}
		}

		if (candidateMatches.size() == 0) {
			return null;
		} else if (candidateMatches.size() == 1) {
			return candidateMatches;
		} else if (candidateMatches.size() > 1) {
			ArrayList<ArrayList<Integer>> segments = getConsistentArrays(candidateMatches);
			// Now pick the top from it

			ArrayList<ArrayList<Integer>> splitSegs = new ArrayList<ArrayList<Integer>>();
			for(int i = 0;i<segments.size();i++) {
				ArrayList<Integer> tmp = segments.get(i);
				ArrayList<ArrayList<Integer>> tmpSegs = splitsArray(tmp);
				for(int j = 0;j<tmpSegs.size();j++) {
					splitSegs.add(tmpSegs.get(j));
				}
			}
			
			double max = 0;
			int maxOp = 0;
			// Actually code to add here, split one sentence group to multiple
			// sentence groups to try every combination
			

			for (int i = 0; i < splitSegs.size(); i++) {
				double sim = getProb(prob, splitSegs.get(i), pair, lineNo,
						aligner, option);
				if (sim > max) {
					max = sim;
					maxOp = i;
				}
			}

			return splitSegs.get(maxOp);
		} else {
			return null;
		}

	}

	public String getAnnStr(ArrayList<Integer> indexes) {
		if (indexes == null)
			return "ADD";
		String ann = "";
		for (int i = 0; i < indexes.size(); i++) {
			ann += (indexes.get(i) + 1) + ",";
		}
		ann = ann.substring(0, ann.length() - 1);
		return ann;
	}

	// The function to annotate the whole directory
	public void auto_annotate(String trainPath, String annotatePath, int option)
			throws Exception {
		LDAligner aligner = new LDAligner();
		Hashtable<DocumentPair, double[][]> annotateVal = aligner.annotate(
				trainPath, annotatePath, option);
		System.err.println("Finished PreProcessing...");
		Iterator<DocumentPair> it = annotateVal.keySet().iterator();
		while (it.hasNext()) {
			DocumentPair pair = it.next();
			double[][] probs = annotateVal.get(pair);
			String filename = pair.getFileName();
			filename = annotatePath + "/" + filename;
			System.err.println("Annotating: " + filename);

			FileInputStream fis = new FileInputStream(new File(filename));
			XSSFWorkbook xwb = new XSSFWorkbook(fis);

			for (int i = 0; i < pair.getModified().getSentences().size(); i++) {
				ArrayList<Integer> alignedIndex = getAlignedSentences(probs, i,
						pair, aligner.getAligner(), option);
				String tag = getAnnStr(alignedIndex);
				writeXSL(xwb, i, tag);
			}

			fis.close();
			FileOutputStream fileOut = new FileOutputStream(filename);
			xwb.write(fileOut);
			fileOut.flush();
			fileOut.close();
		}
	}

	public static void main(String[] args) throws Exception {
		AutomaticAnnotator aa = new AutomaticAnnotator();
		// aa.writeDefault("E:\\independent study\\Revision\\test/ 32bad250-63da-4901-b0f3-3f74f022494a.txt.xlsx");
		String trainPath = "E:\\independent study\\Revision\\train_data\\selected";
		String annotatePath = "E:\\independent study\\Revision\\Braverman\\class1";
		//String annotatePath2 = "E:\\independent study\\Revision\\all_data\\class2";
		int option = 2; // TFIDF
		aa.auto_annotate(trainPath, annotatePath, option);
		//aa.auto_annotate(trainPath, annotatePath2, option);
	}
}
